"""
MCP Solver Utilities

实用函数和辅助类
"""

import torch
import numpy as np
from typing import Dict, Any, List, Tuple, Optional

# 解决Qt平台插件问题 - 使用Agg后端（非交互式）
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt


def check_kkt_conditions(
    controls: torch.Tensor,
    multipliers: torch.Tensor,
    cost_gradients: torch.Tensor,
    constraints: torch.Tensor,
    tolerance: float = 1e-6
) -> Dict[str, float]:
    """
    检查KKT条件

    Args:
        controls: 控制变量
        multipliers: 拉格朗日乘子
        cost_gradients: 成本函数梯度
        constraints: 约束函数值
        tolerance: 容差

    Returns:
        KKT条件违反程度
    """
    # 1. 梯度条件: ∇L = 0
    grad_violation = torch.norm(cost_gradients).item()

    # 2. 原始可行性: g(x) = 0
    primal_feasibility = torch.norm(constraints).item()

    # 3. 对偶可行性: λ ≥ 0
    dual_feasibility = torch.norm(torch.clamp(-multipliers, min=0)).item()

    # 4. 互补松弛: λ^T g(x) = 0
    complementarity = torch.abs(torch.dot(multipliers, constraints)).item()

    return {
        'stationarity': grad_violation,
        'primal_feasibility': primal_feasibility,
        'dual_feasibility': dual_feasibility,
        'complementarity': complementarity,
        'is_kkt_point': all([
            grad_violation < tolerance,
            primal_feasibility < tolerance,
            dual_feasibility < tolerance,
            complementarity < tolerance
        ])
    }


def plot_convergence(history: Dict[str, List[float]], save_path: Optional[str] = None):
    """
    绘制收敛历史

    Args:
        history: 包含损失历史的字典
        save_path: 保存路径
    """
    fig, axes = plt.subplots(2, 2, figsize=(12, 8))

    # 残差
    if 'residual' in history:
        axes[0, 0].semilogy(history['residual'])
        axes[0, 0].set_title('Residual')
        axes[0, 0].set_xlabel('Iteration')
        axes[0, 0].set_ylabel('Residual')
        axes[0, 0].grid(True)

    # 可行性
    if 'feasibility' in history:
        axes[0, 1].semilogy(history['feasibility'])
        axes[0, 1].set_title('Feasibility')
        axes[0, 1].set_xlabel('Iteration')
        axes[0, 1].set_ylabel('Constraint Violation')
        axes[0, 1].grid(True)

    # 互补性
    if 'complementarity' in history:
        axes[1, 0].semilogy(history['complementarity'])
        axes[1, 0].set_title('Complementarity')
        axes[1, 0].set_xlabel('Iteration')
        axes[1, 0].set_ylabel('Complementarity Measure')
        axes[1, 0].grid(True)

    # 平稳性
    if 'stationarity' in history:
        axes[1, 1].semilogy(history['stationarity'])
        axes[1, 1].set_title('Stationarity')
        axes[1, 1].set_xlabel('Iteration')
        axes[1, 1].set_ylabel('Gradient Norm')
        axes[1, 1].grid(True)

    plt.tight_layout()

    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
    plt.close()  # 关闭图形以释放内存


def create_box_constraints(
    n_vars: int,
    lower_bounds: Optional[np.ndarray] = None,
    upper_bounds: Optional[np.ndarray] = None
) -> Tuple[np.ndarray, np.ndarray]:
    """
    创建盒约束

    Args:
        n_vars: 变量数量
        lower_bounds: 下界
        upper_bounds: 上界

    Returns:
        下界和上界数组
    """
    if lower_bounds is None:
        lb = -np.inf * np.ones(n_vars)
    else:
        lb = np.broadcast_to(lower_bounds, n_vars).copy()

    if upper_bounds is None:
        ub = np.inf * np.ones(n_vars)
    else:
        ub = np.broadcast_to(upper_bounds, n_vars).copy()

    return lb, ub


def finite_difference_jacobian(
    func: callable,
    x: np.ndarray,
    eps: float = 1e-8
) -> np.ndarray:
    """
    计算有限差分雅可比矩阵

    Args:
        func: 函数
        x: 点
        eps: 扰动大小

    Returns:
        雅可比矩阵
    """
    n = len(x)
    f0 = func(x)
    m = len(f0)

    J = np.zeros((m, n))

    for j in range(n):
        x_plus = x.copy()
        x_plus[j] += eps

        x_minus = x.copy()
        x_minus[j] -= eps

        f_plus = func(x_plus)
        f_minus = func(x_minus)

        J[:, j] = (f_plus - f_minus) / (2 * eps)

    return J


def analyze_mcp_solution(
    solution: Dict[str, Any],
    F_func: callable,
    tolerance: float = 1e-6
) -> Dict[str, Any]:
    """
    分析MCP解的质量

    Args:
        solution: MCP求解结果
        F_func: MCP函数
        tolerance: 容差

    Returns:
        分析结果
    """
    z = solution['solution']
    F_val = F_func(z)

    analysis = {
        'is_solution': solution.get('success', False),
        'residual_norm': np.linalg.norm(F_val),
        'max_residual': np.max(np.abs(F_val)),
        'satisfies_tolerance': np.linalg.norm(F_val) < tolerance,
    }

    # 如果有边界信息
    if 'lb' in solution and 'ub' in solution:
        lb, ub = solution['lb'], solution['ub']

        # 检查变量边界
        lb_violations = np.sum(z < lb - tolerance)
        ub_violations = np.sum(z > ub + tolerance)

        analysis.update({
            'lb_violations': lb_violations,
            'ub_violations': ub_violations,
            'satisfies_bounds': lb_violations == 0 and ub_violations == 0
        })

    return analysis


class MCPMonitor:
    """
    MCP求解监控器

    记录求解过程中的统计信息
    """

    def __init__(self):
        self.history = {
            'iteration': [],
            'residual': [],
            'stationarity': [],
            'feasibility': [],
            'complementarity': [],
            'step_size': [],
            'function_evals': 0,
            'jacobian_evals': 0
        }

    def record(self, iteration: int, **metrics):
        """记录一次迭代的指标"""
        self.history['iteration'].append(iteration)

        for key, value in metrics.items():
            if key in self.history:
                self.history[key].append(value)

    def increment_evals(self, func_evals: int = 0, jac_evals: int = 0):
        """增加函数和雅可比评估次数"""
        self.history['function_evals'] += func_evals
        self.history['jacobian_evals'] += jac_evals

    def plot(self, save_path: Optional[str] = None):
        """绘制收敛历史"""
        plot_convergence(self.history, save_path)

    def get_summary(self) -> Dict[str, Any]:
        """获取求解摘要"""
        if not self.history['iteration']:
            return {'status': 'no_data'}

        return {
            'total_iterations': len(self.history['iteration']),
            'final_residual': self.history['residual'][-1] if self.history['residual'] else None,
            'function_evaluations': self.history['function_evals'],
            'jacobian_evaluations': self.history['jacobian_evals'],
            'converged': (self.history['residual'][-1] < 1e-6) if self.history['residual'] else False
        }


def warm_start_from_previous(
    previous_solution: torch.Tensor,
    new_problem_size: int,
    strategy: str = 'pad_zeros'
) -> torch.Tensor:
    """
    从上一个解生成热启动点

    Args:
        previous_solution: 上一个问题的解
        new_problem_size: 新问题的变量数
        strategy: 策略 ('pad_zeros', 'repeat_last', 'interpolate')

    Returns:
        热启动点
    """
    prev_size = len(previous_solution)

    if prev_size == new_problem_size:
        return previous_solution.clone()

    elif prev_size < new_problem_size:
        # 需要扩展
        if strategy == 'pad_zeros':
            warm_start = torch.zeros(new_problem_size)
            warm_start[:prev_size] = previous_solution
        elif strategy == 'repeat_last':
            warm_start = torch.zeros(new_problem_size)
            warm_start[:prev_size] = previous_solution
            warm_start[prev_size:] = previous_solution[-1]
        else:  # interpolate
            warm_start = torch.zeros(new_problem_size)
            warm_start[:prev_size] = previous_solution
            # 简单线性插值
            for i in range(prev_size, new_problem_size):
                alpha = (i - prev_size) / (new_problem_size - prev_size)
                warm_start[i] = (1 - alpha) * previous_solution[-1]

    else:
        # 需要截断
        warm_start = previous_solution[:new_problem_size]

    return warm_start


def validate_mcp_inputs(
    F_func: callable,
    z0: np.ndarray,
    lb: Optional[np.ndarray] = None,
    ub: Optional[np.ndarray] = None
) -> bool:
    """
    验证MCP输入的有效性

    Args:
        F_func: MCP函数
        z0: 初始点
        lb: 下界
        ub: 上界

    Returns:
        输入是否有效
    """
    try:
        # 检查初始点
        if not isinstance(z0, np.ndarray) or z0.ndim != 1:
            print("错误: z0必须是一维numpy数组")
            return False

        n = len(z0)

        # 检查函数
        F_val = F_func(z0)
        if not isinstance(F_val, np.ndarray) or F_val.ndim != 1:
            print("错误: F函数必须返回一维numpy数组")
            return False

        if len(F_val) != n:
            print(f"错误: F函数输出维度({len(F_val)})与输入维度({n})不匹配")
            return False

        # 检查边界
        if lb is not None:
            if len(lb) != n:
                print(f"错误: 下界维度({len(lb)})与变量维度({n})不匹配")
                return False

            if np.any(z0 < lb):
                print("警告: 初始点违反下界约束")

        if ub is not None:
            if len(ub) != n:
                print(f"错误: 上界维度({len(ub)})与变量维度({n})不匹配")
                return False

            if np.any(z0 > ub):
                print("警告: 初始点违反上界约束")

        if lb is not None and ub is not None:
            if np.any(lb > ub):
                print("错误: 存在下界大于上界的情况")
                return False

        return True

    except Exception as e:
        print(f"输入验证出错: {e}")
        return False